//
// Created by DELL on 2022/4/22.
//

#include "DescriptorSet.h"
#include "Device.h"
#include "UniformBuffer.h"
#include "Texture.h"

BEGIN_NS_VULKAN

VkDescriptorSetLayout DescriptorSetLayout::createVkLayout(VkDevice device)
{
    std::vector<VkDescriptorSetLayoutBinding> bindings;
    bindings.resize(this->samples.size() + this->blocks.size());
    int i = 0;
    for (auto& block : blocks)
    {
        auto& uboLayoutBinding = bindings[i];
        uboLayoutBinding.binding = block.binding;
        uboLayoutBinding.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
        uboLayoutBinding.descriptorCount = 1;
        uboLayoutBinding.stageFlags = block.stages;
        uboLayoutBinding.pImmutableSamplers = nullptr; // Optional
        ++i;
    }
    for (auto& sampler : samples)
    {
        auto& samplerLayoutBinding = bindings[i];;
        samplerLayoutBinding.binding = sampler.second.binging;
        samplerLayoutBinding.descriptorCount = 1;
        samplerLayoutBinding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
        samplerLayoutBinding.pImmutableSamplers = nullptr;
        samplerLayoutBinding.stageFlags =  sampler.second.stages;
        ++i;
    }

    VkDescriptorSetLayoutCreateInfo layoutInfo{};
    layoutInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
    layoutInfo.bindingCount = static_cast<uint32_t>(bindings.size());
    layoutInfo.pBindings = bindings.data();

    VkDescriptorSetLayout layout{};
    if (vkCreateDescriptorSetLayout(device, &layoutInfo, nullptr, &layout) != VK_SUCCESS) {
        throw std::runtime_error("failed to create descriptor set layout!");
    }
    return layout;
}

DescriptorSet::DescriptorSet(DeviceRef device)
    : device(device)
{

}

DescriptorSet::~DescriptorSet()
{
//    vkFreeDescriptorSets(device->getDevice(), descriptorPool, descriptorSets.size(), descriptorSets.data());
}

DescriptorSetRef DescriptorSet::create(DeviceRef& device, VkDescriptorSetLayout layout, VkDescriptorPool pool)
{
    auto descriptor = new DescriptorSet(device);

    VkDescriptorSetAllocateInfo allocInfo{};
    allocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
    allocInfo.descriptorPool = pool;
    allocInfo.descriptorSetCount = 1;
    allocInfo.pSetLayouts = &layout;

    if (vkAllocateDescriptorSets(device->getDevice(), &allocInfo, &descriptor->descriptorSet) != VK_SUCCESS) {
        throw std::runtime_error("failed to allocate descriptor sets!");
    }
    descriptor->bufferInfos.reserve(4);
    descriptor->imageInfos.reserve(16);

    return Vulkan::DescriptorSetRef(descriptor);
}

void DescriptorSet::addBuffer(UniformBufferRef& uniform, int binding)
{
    this->bufferInfos.push_back({});
    auto& bufferInfo = bufferInfos.back();
    bufferInfo.buffer = uniform->getBuffer();
    bufferInfo.offset = 0;
    bufferInfo.range = uniform->getSize();

    this->writeDescriptorSet.push_back({});
    auto& descriptorWrite = this->writeDescriptorSet.back();
    descriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    descriptorWrite.dstSet = descriptorSet;
    descriptorWrite.dstBinding = binding;
    descriptorWrite.dstArrayElement = 0;
    descriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
    descriptorWrite.descriptorCount = 1;
    descriptorWrite.pBufferInfo = &bufferInfo;

}

void DescriptorSet::addSampler(TextureRef& texture, VkSampler sampler, int binding)
{
    this->addSampler(texture->getImageView(), sampler, binding);
}

void DescriptorSet::addSampler(VkImageView imageView, VkSampler sampler, int binding)
{
    this->imageInfos.push_back({});
    auto& imageInfo = this->imageInfos.back();
    imageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    imageInfo.imageView = imageView;
    imageInfo.sampler = sampler;

    this->writeDescriptorSet.push_back({});
    auto& descriptorWrite = this->writeDescriptorSet.back();

    descriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    descriptorWrite.dstSet = descriptorSet;
    descriptorWrite.dstBinding = binding;
    descriptorWrite.dstArrayElement = 0;
    descriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    descriptorWrite.descriptorCount = 1;
    descriptorWrite.pImageInfo = &imageInfo;
}

void DescriptorSet::build()
{
    if(this->writeDescriptorSet.empty())
    {
        return;
    }

    vkUpdateDescriptorSets(device->getDevice(), static_cast<uint32_t>(writeDescriptorSet.size()),
                           writeDescriptorSet.data(), 0, nullptr);

    this->bufferInfos.clear();
    this->writeDescriptorSet.clear();
    this->imageInfos.clear();

}

END_NS_VULKAN
